/*
 * Copyright 2025 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.google.firebase.ai

import com.google.firebase.FirebaseApp
import com.google.firebase.ai.common.APIController
import com.google.firebase.ai.common.AppCheckHeaderProvider
import com.google.firebase.ai.common.TemplateGenerateContentRequest
import com.google.firebase.ai.type.Content
import com.google.firebase.ai.type.FinishReason
import com.google.firebase.ai.type.FirebaseAIException
import com.google.firebase.ai.type.GenerateContentResponse
import com.google.firebase.ai.type.PromptBlockedException
import com.google.firebase.ai.type.PublicPreviewAPI
import com.google.firebase.ai.type.RequestOptions
import com.google.firebase.ai.type.ResponseStoppedException
import com.google.firebase.ai.type.SerializationException
import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
import com.google.firebase.auth.internal.InternalAuthProvider
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.catch
import kotlinx.coroutines.flow.map
import kotlinx.serialization.json.Json
import kotlinx.serialization.json.jsonObject
import org.json.JSONObject

/**
 * Represents a multimodal model (like Gemini), capable of generating content based on various
 * templated input types.
 */
@PublicPreviewAPI
public class TemplateGenerativeModel
internal constructor(
  private val templateUri: String,
  private val controller: APIController,
) {

  internal constructor(
    templateUri: String,
    apiKey: String,
    firebaseApp: FirebaseApp,
    useLimitedUseAppCheckTokens: Boolean,
    requestOptions: RequestOptions = RequestOptions(),
    appCheckTokenProvider: InteropAppCheckTokenProvider? = null,
    internalAuthProvider: InternalAuthProvider? = null
  ) : this(
    templateUri,
    APIController(
      apiKey,
      "",
      requestOptions,
      "gl-kotlin/${KotlinVersion.CURRENT}-ai fire/${BuildConfig.VERSION_NAME}",
      firebaseApp,
      AppCheckHeaderProvider(
        TAG,
        useLimitedUseAppCheckTokens,
        appCheckTokenProvider,
        internalAuthProvider
      ),
    ),
  )

  /**
   * Generates content from a prompt template and inputs.
   *
   * @param templateId The ID of the prompt template to use.
   * @param inputs A map of variables to substitute into the template.
   * @return The content generated by the model.
   * @throws [FirebaseAIException] if the request failed.
   * @see [FirebaseAIException] for types of errors.
   */
  public suspend fun generateContent(
    templateId: String,
    inputs: Map<String, Any>,
  ): GenerateContentResponse =
    try {
      controller
        .templateGenerateContent("$templateUri$templateId", constructRequest(inputs))
        .toPublic()
        .validate()
    } catch (e: Throwable) {
      throw FirebaseAIException.from(e)
    }

  /**
   * Generates content as a stream from a prompt template and inputs.
   *
   * @param templateId The ID of the prompt template to use.
   * @param inputs A map of variables to substitute into the template.
   * @return A [Flow] which will emit responses as they are returned by the model.
   * @throws [FirebaseAIException] if the request failed.
   * @see [FirebaseAIException] for types of errors.
   */
  public fun generateContentStream(
    templateId: String,
    inputs: Map<String, Any>
  ): Flow<GenerateContentResponse> =
    controller
      .templateGenerateContentStream("$templateUri$templateId", constructRequest(inputs))
      .catch { throw FirebaseAIException.from(it) }
      .map { it.toPublic().validate() }

  internal fun constructRequest(
    inputs: Map<String, Any>,
    history: List<Content>? = null
  ): TemplateGenerateContentRequest {
    return TemplateGenerateContentRequest(
      Json.parseToJsonElement(JSONObject(inputs).toString()).jsonObject,
      history?.let { it.map { it.toTemplateInternal() } }
    )
  }

  private fun GenerateContentResponse.validate() = apply {
    if (candidates.isEmpty() && promptFeedback == null) {
      throw SerializationException("Error deserializing response, found no valid fields")
    }
    promptFeedback?.blockReason?.let { throw PromptBlockedException(this) }
    candidates
      .mapNotNull { it.finishReason }
      .firstOrNull { it != FinishReason.STOP }
      ?.let { throw ResponseStoppedException(this) }
  }

  private companion object {
    private val TAG = TemplateGenerativeModel::class.java.simpleName
  }
}
